{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "using RigidBodyDynamics\n",
    "using StaticArrays\n",
    "using SymPy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create symbolic parameters\n",
    "* Masses: $m_1, m_2$\n",
    "* Mass moments of inertia (about center of mass): $I_1, I_2$\n",
    "* Link lengths: $l_1, l_2$\n",
    "* Center of mass locations (w.r.t. preceding joint axis): $c_1, c_2$\n",
    "* Gravitational acceleration: $g$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "\\begin{bmatrix}m_{1}&m_{2}&I_{1}&I_{2}&l_{1}&l_{2}&c_{1}&c_{2}&g\\end{bmatrix}"
      ],
      "text/plain": [
       "1×9 RowVector{SymPy.Sym,Array{SymPy.Sym,1}}:\n",
       " m_1  m_2  I_1  I_2  l_1  l_2  c_1  c_2  g"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inertias = @syms m_1 m_2 I_1 I_2 positive = true\n",
    "lengths = @syms l_1 l_2 c_1 c_2 real = true\n",
    "gravitational_acceleration = @syms g real = true\n",
    "params = [inertias..., lengths..., gravitational_acceleration...]\n",
    "params.'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create double pendulum `Mechanism`\n",
    "\n",
    "A `Mechanism` contains the joint layout and inertia parameters, but no state information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Spanning tree:\n",
       "Vertex: world (root)\n",
       "  Vertex: upper_link, Edge: shoulder\n",
       "    Vertex: lower_link, Edge: elbow\n",
       "No non-tree joints."
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "T = Sym # the 'scalar type' of the Mechanism we'll construct\n",
    "axis = SVector(zero(T), one(T), zero(T)) # axis of rotation for each of the joints\n",
    "double_pendulum = Mechanism(RigidBody{T}(\"world\"); gravity = SVector(zero(T), zero(T), g))\n",
    "world = root_body(double_pendulum) # the fixed 'world' rigid body\n",
    "\n",
    "# Attach the first (upper) link to the world via a revolute joint named 'shoulder'\n",
    "inertia1 = SpatialInertia(CartesianFrame3D(\"upper_link\"), I_1 * axis * axis.', m_1 * SVector(zero(T), zero(T), c_1), m_1)\n",
    "body1 = RigidBody(inertia1)\n",
    "joint1 = Joint(\"shoulder\", Revolute(axis))\n",
    "joint1_to_world = eye(Transform3D{T}, frame_before(joint1), default_frame(world))\n",
    "attach!(double_pendulum, world, body1, joint1, joint_pose = joint1_to_world)\n",
    "\n",
    "# Attach the second (lower) link to the world via a revolute joint named 'elbow'\n",
    "inertia2 = SpatialInertia(CartesianFrame3D(\"lower_link\"), I_2 * axis * axis.', m_2 * SVector(zero(T), zero(T), c_2), m_2)\n",
    "body2 = RigidBody(inertia2)\n",
    "joint2 = Joint(\"elbow\", Revolute(axis))\n",
    "joint2_to_body1 = Transform3D(frame_before(joint2), default_frame(body1), SVector(zero(T), zero(T), l_1))\n",
    "attach!(double_pendulum, body1, body2, joint2, joint_pose = joint2_to_body1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create `MechanismState` associated with the double pendulum `Mechanism`\n",
    "\n",
    "A `MechanismState` stores all state-dependent information associated with a `Mechanism`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x = MechanismState(double_pendulum);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "\\begin{bmatrix}q_{1}\\\\q_{2}\\end{bmatrix}"
      ],
      "text/plain": [
       "2-element Array{SymPy.Sym,1}:\n",
       " q_1\n",
       " q_2"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Set the joint configuration vector of the MechanismState to a new vector of symbolic variables\n",
    "configuration(x)[:] = [symbols(\"q_$i\", real = true) for i = 1 : num_positions(x)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "\\begin{bmatrix}v_{1}\\\\v_{2}\\end{bmatrix}"
      ],
      "text/plain": [
       "2-element Array{SymPy.Sym,1}:\n",
       " v_1\n",
       " v_2"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Set the joint velocity vector of the MechanismState to a new vector of symbolic variables\n",
    "velocity(x)[:] = [symbols(\"v_$i\", real = true) for i = 1 : num_positions(x)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute dynamical quantities in symbolic form"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "\\begin{bmatrix}I_{1} + I_{2} + 2 c_{2} l_{1} m_{2} \\cos{\\left (q_{2} \\right )} + l_{1}^{2} m_{2}&I_{2} + c_{2} l_{1} m_{2} \\cos{\\left (q_{2} \\right )}\\\\I_{2} + c_{2} l_{1} m_{2} \\cos{\\left (q_{2} \\right )}&I_{2}\\end{bmatrix}"
      ],
      "text/plain": [
       "2×2 Symmetric{SymPy.Sym,Array{SymPy.Sym,2}}:\n",
       " I_1 + I_2 + 2*c_2*l_1*m_2*cos(q_2) + l_1^2*m_2  I_2 + c_2*l_1*m_2*cos(q_2)\n",
       "                     I_2 + c_2*l_1*m_2*cos(q_2)                         I_2"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Mass matrix\n",
    "M = mass_matrix(x)\n",
    "map!(simplify, M.data, M.data) # Note: M is a Symmetric matrix type; need to simplify the underlying data\n",
    "M"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false,
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$$\\frac{I_{1} v_{1}^{2}}{2} + \\frac{I_{2} v_{1}^{2}}{2} + I_{2} v_{1} v_{2} + \\frac{I_{2} v_{2}^{2}}{2} + c_{2} l_{1} m_{2} v_{1}^{2} \\cos{\\left (q_{2} \\right )} + c_{2} l_{1} m_{2} v_{1} v_{2} \\cos{\\left (q_{2} \\right )} + \\frac{l_{1}^{2} m_{2}}{2} v_{1}^{2}$$"
      ],
      "text/plain": [
       "       2          2                        2                                  \n",
       "I_1*v_1    I_2*v_1                  I_2*v_2                   2               \n",
       "-------- + -------- + I_2*v_1*v_2 + -------- + c_2*l_1*m_2*v_1 *cos(q_2) + c_2\n",
       "   2          2                        2                                      \n",
       "\n",
       "                               2        2\n",
       "                            l_1 *m_2*v_1 \n",
       "*l_1*m_2*v_1*v_2*cos(q_2) + -------------\n",
       "                                  2      "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Kinetic energy\n",
    "simplify(kinetic_energy(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$$- g \\left(c_{1} m_{1} \\cos{\\left (q_{1} \\right )} + c_{2} m_{2} \\cos{\\left (q_{1} + q_{2} \\right )} + l_{1} m_{2} \\cos{\\left (q_{1} \\right )}\\right)$$"
      ],
      "text/plain": [
       "-g*(c_1*m_1*cos(q_1) + c_2*m_2*cos(q_1 + q_2) + l_1*m_2*cos(q_1))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Potential energy\n",
    "simplify(gravitational_potential_energy(x))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 0.6.2",
   "language": "julia",
   "name": "julia-0.6"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "0.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
